use std::collections::{HashMap, HashSet};

use crate::{
    contracts_items::{
        get_invariant_expl, is_assertion, is_before_loop, is_loop_variant, is_snapshot_closure,
    },
    ctx::{HasTyCtxt, TranslationCtx},
    translation::pearlite::Term,
};
use rustc_hir::def_id::DefId;
use rustc_middle::{
    mir::{AggregateKind, BasicBlock, Body, Location, Operand, Rvalue, visit::Visitor},
    ty::{Ty, TyCtxt, TyKind},
};

#[derive(Debug, Clone)]
pub enum LoopSpecKind {
    Invariant(String),
    Variant,
}

/// Collect spec-related special closures in a mir [`Body`].
pub(crate) struct SpecClosures<'tcx> {
    /// Closures generated by `proof_assert!`
    pub(crate) assertions: HashMap<DefId, Term<'tcx>>,
    /// Closures generated by `snapshot!`
    pub(crate) snapshots: HashMap<DefId, Term<'tcx>>,
}

impl<'tcx> SpecClosures<'tcx> {
    pub(crate) fn collect(ctx: &TranslationCtx<'tcx>, body: &Body<'tcx>) -> Self {
        let mut visitor = Closures::new(ctx.tcx);
        visitor.visit_body(body);

        let mut assertions = HashMap::new();
        let mut snapshots = HashMap::new();
        for clos in visitor.closures.into_iter() {
            if is_assertion(ctx.tcx, clos) {
                let term = ctx.term(clos).unwrap().1.clone();
                assertions.insert(clos, term);
            } else if is_snapshot_closure(ctx.tcx, clos) {
                let term = ctx.term(clos).unwrap().1.clone();
                snapshots.insert(clos, term);
            }
        }
        Self { assertions, snapshots }
    }
}

// Collect the closures in thir, so that we can do typechecking ourselves, and
// translate the invariant closure from thir.
struct Closures<'tcx> {
    pub tcx: TyCtxt<'tcx>,
    pub closures: HashSet<DefId>,
}

impl<'tcx> Closures<'tcx> {
    fn new(tcx: TyCtxt<'tcx>) -> Self {
        Closures { tcx, closures: HashSet::new() }
    }
}

fn snapshot_closure_id<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Option<DefId> {
    if let TyKind::Closure(def_id, _) = ty.peel_refs().kind() {
        is_snapshot_closure(tcx, *def_id).then_some(*def_id)
    } else {
        None
    }
}

impl<'tcx> Visitor<'tcx> for Closures<'tcx> {
    fn visit_rvalue(&mut self, rvalue: &Rvalue<'tcx>, loc: Location) {
        match rvalue {
            Rvalue::Aggregate(box AggregateKind::Closure(id, _), _) => {
                self.closures.insert(*id);
            }
            Rvalue::Use(Operand::Constant(box ck)) => {
                if let Some(def_id) = snapshot_closure_id(self.tcx, ck.const_.ty()) {
                    self.closures.insert(def_id);
                }
            }
            _ => {}
        }
        self.super_rvalue(rvalue, loc);
    }
}

pub(crate) struct InvariantsAndVariants<'tcx> {
    /// Invariants placed at the beginning of their respective loops.
    ///
    /// The string is a description for Why3.
    pub(crate) invariants: HashMap<BasicBlock, Vec<(String, Term<'tcx>)>>,
    /// Variants placed at the beginning of their respective loops.
    pub(crate) variants: HashMap<BasicBlock, Term<'tcx>>,
    /// Invariants for which we couldn't find a loop header are translated as assertions.
    pub(crate) assertions: HashMap<DefId, (Term<'tcx>, String)>,
}

struct InvariantsVisitor<'a, 'tcx> {
    ctx: &'a TranslationCtx<'tcx>,
    body: &'a Body<'tcx>,
    before_loop: HashSet<BasicBlock>,
    invariants_and_variants: InvariantsAndVariants<'tcx>,
}

impl InvariantsVisitor<'_, '_> {
    /// Detect if a loop really loops.
    ///
    /// This searches for a loop header: a special marker inserted just before
    /// a loop.
    ///
    /// `loc` should be the MIR location of an (in)variant: that is, it should be
    /// right after the loop entry.
    ///
    /// If the loop really loops, its entry should have more than one
    /// predecessor: so we search backward for such a block.
    fn find_loop_header(&self, loc: Location) -> Option<BasicBlock> {
        let mut block = loc.block;
        if self.before_loop.contains(&block) {
            // Reached "before_loop" marker in the same block.
            // This assumes that statements are visited in order, so that if a block
            // contains both invariants and a "before_block" marker, the marker is not
            // in the `before_loop` set when we visit invariants before it.
            return None;
        }
        loop {
            let preds = &self.body.basic_blocks.predecessors()[block];
            if preds.len() > 1 {
                return Some(block);
            }
            let Some(pred) = preds.first() else {
                // Reached the top of the function. Impossible.
                panic!("The impossible happened: Missing 'before_loop' marker.");
            };
            if self.before_loop.contains(pred) {
                // Reached "before_loop" marker.
                return None;
            }
            block = *pred;
        }
    }
}

impl<'tcx> Visitor<'tcx> for InvariantsVisitor<'_, 'tcx> {
    fn visit_rvalue(&mut self, rvalue: &Rvalue<'tcx>, loc: Location) {
        if let Rvalue::Aggregate(box AggregateKind::Closure(id, _), _) = rvalue {
            let kind = if let Some(expl) = get_invariant_expl(self.ctx.tcx, *id) {
                LoopSpecKind::Invariant(expl)
            } else if is_loop_variant(self.ctx.tcx, *id) {
                LoopSpecKind::Variant
            } else {
                if is_before_loop(self.ctx.tcx, *id) {
                    self.before_loop.insert(loc.block);
                }
                return;
            };
            let term = self.ctx.term(*id).unwrap().1.clone();
            match self.find_loop_header(loc) {
                None if let LoopSpecKind::Invariant(expl) = kind => {
                    self.ctx.warn(
                        self.ctx.def_span(id),
                        "This loop does not loop. This invariant could just be an assertion.",
                    );
                    self.invariants_and_variants.assertions.insert(*id, (term, expl));
                }
                None => self.ctx.warn(
                    self.ctx.def_span(id),
                    "This loop does not loop. This variant will be ignored.",
                ),
                Some(target) => match kind {
                    LoopSpecKind::Invariant(desc) => self
                        .invariants_and_variants
                        .invariants
                        .entry(target)
                        .or_default()
                        .push((desc, term)),
                    LoopSpecKind::Variant => {
                        if self.invariants_and_variants.variants.insert(target, term).is_some() {
                            self.ctx.crash_and_error(
                                self.body.span,
                                "Only one variant can be provided for each loop",
                            );
                        }
                    }
                },
            }
        }
        self.super_rvalue(rvalue, loc);
    }
}

impl<'a, 'tcx> HasTyCtxt<'tcx> for InvariantsVisitor<'a, 'tcx> {
    fn tcx(&self) -> TyCtxt<'tcx> {
        self.ctx.tcx
    }
}

/// Calculate the *actual* location of invariants in MIR
pub(crate) fn corrected_invariant_names_and_locations<'tcx>(
    ctx: &TranslationCtx<'tcx>,
    body: &Body<'tcx>,
) -> InvariantsAndVariants<'tcx> {
    let mut invs_gather = InvariantsVisitor {
        ctx,
        body,
        before_loop: HashSet::new(),
        invariants_and_variants: InvariantsAndVariants {
            invariants: HashMap::new(),
            variants: HashMap::new(),
            assertions: HashMap::new(),
        },
    };
    invs_gather.visit_body(body);
    invs_gather.invariants_and_variants
}
